// FSDrives.cpp

#include "StdAfx.h"

#include "../C/Alloc.h"

#include "Common/String/StringConvert.h"

#include "Common/System/Defs.h"
#include "FileDir.h"
#include "Common/Stream/FileIO.h"
#include "FileSystem.h"
#include "Common/MS/PropVariant.h"

#include "Common/7Zip/PropID.h"

#include "FSDrives.h"
#include "FSFolder.h"
#include "Common/Lang/LangUtils.h"
#include "SysIconUtils.h"

//TODO: AM - this is very bad
#include "../../7zip/UI/FileManager/resource.h"

#include <algorithm>

#include "Common/Defs.h"

using namespace Common;


namespace Common
{
namespace FileSystem
{

static LPCTSTR kVolPrefix = _T("\\\\.\\");

CString CDriveInfo::GetDeviceFileIoName() const
{
	return kVolPrefix + Name;
}

struct CPhysTempBuffer
{
	void* buffer;
	CPhysTempBuffer(): buffer(0) {}
	~CPhysTempBuffer()
	{
		MidFree(buffer);
	}
};

static HRESULT CopyFileSpec(LPCWSTR fromPath, LPCWSTR toPath, bool writeToDisk, UINT64 fileSize,
							UINT32 bufferSize, UINT64 progressStart, IProgress* progress)
{
	Common::IO::CInFile inFile;
	if (!inFile.Open(fromPath))
	{
		return GetLastError();
	}
	if (fileSize == (UINT64)(INT64) - 1)
	{
		if (!inFile.GetLength(fileSize))
		{
			::GetLastError();
		}
	}
	Common::IO::COutFile outFile;
	if (writeToDisk)
	{
		if (!outFile.Open(toPath, FILE_SHARE_WRITE, OPEN_EXISTING, 0))
		{
			return GetLastError();
		}
	}
	else if (!outFile.Create(toPath, true))
	{
		return GetLastError();
	}
	CPhysTempBuffer tempBuffer;
	tempBuffer.buffer = MidAlloc(bufferSize);
	if (tempBuffer.buffer == 0)
	{
		return E_OUTOFMEMORY;
	}

	HRESULT hRes (S_OK);
	for (UINT64 pos = 0; pos < fileSize;)
	{
		UINT64 progressCur = progressStart + pos;
		hRes = progress->SetCompleted(&progressCur);
		if(S_OK != hRes)
		{
			return hRes;
		}
		UINT64 rem = fileSize - pos;
		UINT32 curSize = (UINT32)min(rem, (UINT64)bufferSize);
		UINT32 processedSize;
		if (!inFile.Read(tempBuffer.buffer, curSize, processedSize))
		{
			return GetLastError();
		}
		if (processedSize == 0)
		{
			break;
		}
		curSize = processedSize;
		if (writeToDisk)
		{
			const UINT32 kMask = 0x1FF;
			curSize = (curSize + kMask) & ~kMask;
			if (curSize > bufferSize)
			{
				return E_FAIL;
			}
		}

		if (!outFile.Write(tempBuffer.buffer, curSize, processedSize))
		{
			return GetLastError();
		}
		if (curSize != processedSize)
		{
			return E_FAIL;
		}
		pos += curSize;
	}
	return S_OK;
}

static const STATPROPSTG kProps[] =
{
	{ NULL, kpidName, VT_BSTR},
	{ NULL, kpidTotalSize, VT_UI8},
	{ NULL, kpidFreeSpace, VT_UI8},
	{ NULL, kpidType, VT_BSTR},
	{ NULL, kpidVolumeName, VT_BSTR},
	{ NULL, kpidFileSystem, VT_BSTR},
	{ NULL, kpidClusterSize, VT_UI8}
};

static const char* kDriveTypes[] =
{
	"Unknown",
	"No Root Dir",
	"Removable",
	"Fixed",
	"Remote",
	"CD-ROM",
	"RAM disk"
};

STDMETHODIMP CFSDrives::LoadItems()
{
	_drives.clear();

	CStringVector driveStrings;
	FileSystem::Find::MyGetLogicalDriveStrings(driveStrings);
	_drives.reserve(driveStrings.size());
	for (size_t i = 0; i < driveStrings.size(); i++)
	{
		CDriveInfo di;

		const CString &driveName = driveStrings[i];

		di.FullSystemName = driveName;

		di.Name = di.FullSystemName.Left(di.FullSystemName.GetLength() - 1);
		di.ClusterSize = 0;
		di.DriveSize = 0;
		di.FreeSpace = 0;
		di.DriveType = ::GetDriveType(driveName);
		bool needRead = true;
		if (di.DriveType == DRIVE_CDROM || di.DriveType == DRIVE_REMOVABLE)
		{
			/*
			DWORD dwSerialNumber;`
			if (!::GetVolumeInformation(di.FullSystemName,
			    NULL, 0, &dwSerialNumber, NULL, NULL, NULL, 0))
			*/
			di.KnownSizes = false;
			{
				needRead = false;
			}
		}
		if (needRead)
		{
			TCHAR volumeName[MAX_PATH + 1] = {0};
			TCHAR fileSystemName[MAX_PATH + 1] = {0};
			DWORD volumeSerialNumber, maximumComponentLength, fileSystemFlags;
			::GetVolumeInformation(driveName,
								   volumeName, MAX_PATH,
								   &volumeSerialNumber, &maximumComponentLength, &fileSystemFlags,
								   fileSystemName, MAX_PATH);
			di.VolumeName = volumeName;
			di.FileSystemName = fileSystemName;

			FileSystem::GetDiskFreeSpace(driveName,
										 di.ClusterSize, di.DriveSize, di.FreeSpace);
			di.KnownSizes = true;
		}
		_drives.push_back(di);
	}
	return S_OK;
}

STDMETHODIMP CFSDrives::GetNumberOfItems(UINT32* numItems)
{
	*numItems = _drives.size();
	return S_OK;
}

STDMETHODIMP CFSDrives::GetProperty(UINT32 itemIndex, PROPID propID, PROPVARIANT* value)
{
	if (itemIndex >= (UINT32)_drives.size())
	{
		return E_INVALIDARG;
	}
	CComPropVariant prop;
	const CDriveInfo &di = _drives[itemIndex];
	switch(propID)
	{
		case kpidIsDir:
			prop = !_volumeMode;
			break;
		case kpidName:
			prop = di.Name;
			break;
		case kpidTotalSize:
			if (di.KnownSizes)
			{
				prop = di.DriveSize;
			}
			break;
		case kpidFreeSpace:
			if (di.KnownSizes)
			{
				prop = di.FreeSpace;
			}
			break;
		case kpidClusterSize:
			if (di.KnownSizes)
			{
				prop = di.ClusterSize;
			}
			break;
		case kpidType:
			if (di.DriveType < sizeof(kDriveTypes) / sizeof(kDriveTypes[0]))
			{
				prop = kDriveTypes[di.DriveType];
			}
			break;
		case kpidVolumeName:
			prop =  CComBSTR(di.VolumeName);
			break;
		case kpidFileSystem:
			prop = CComBSTR(di.FileSystemName);
			break;
	}
	prop.Detach(value);
	return S_OK;
}

HRESULT CFSDrives::BindToFolderSpec(const wchar_t* name, IFolderFolder** resultFolder)
{
	*resultFolder = 0;
	if (_volumeMode)
	{
		return S_OK;
	}
	FileSystem::Folder::CFSFolder* fsFolderSpec = new CComObject<FileSystem::Folder::CFSFolder>;
	fsFolderSpec->FinalConstruct();
	CComPtr<IFolderFolder> subFolder = fsFolderSpec;
	HRESULT hRes(fsFolderSpec->Init(name, 0));
	if(S_OK != hRes)
	{
		return hRes;
	}
	*resultFolder = subFolder.Detach();
	return S_OK;
}

STDMETHODIMP CFSDrives::BindToFolder(UINT32 index, IFolderFolder** resultFolder)
{
	*resultFolder = 0;
	if (index >= (UINT32)_drives.size())
	{
		return E_INVALIDARG;
	}
	const CDriveInfo &di = _drives[index];
	/*
	if (_volumeMode)
	{
	  *resultFolder = 0;
	  CPhysDriveFolder *folderSpec = new CPhysDriveFolder;
	  CComPtr<IFolderFolder> subFolder = folderSpec;
	  RINOK(folderSpec->Init(di.Name));
	  *resultFolder = subFolder.Detach();
	  return S_OK;
	}
	*/
	return BindToFolderSpec(di.FullSystemName, resultFolder);
}

STDMETHODIMP CFSDrives::BindToFolder(const wchar_t* name, IFolderFolder** resultFolder)
{
	return BindToFolderSpec(name, resultFolder);
}

STDMETHODIMP CFSDrives::BindToParentFolder(IFolderFolder** resultFolder)
{
	*resultFolder = 0;
	return S_OK;
}

STDMETHODIMP CFSDrives::GetNumberOfProperties(UINT32* numProperties)
{
	*numProperties = sizeof(kProps) / sizeof(kProps[0]);
	return S_OK;
}
STDMETHODIMP CFSDrives::GetPropertyInfo(UINT32 index, BSTR* name, PROPID* propID, VARTYPE* varType)
{
	if(index >= sizeof(kProps) / sizeof(kProps[0]))
	{
		return E_INVALIDARG;
	}
	const STATPROPSTG &srcItem = kProps[index];
	*propID = srcItem.propid;
	*varType = srcItem.vt;
	*name = 0;
	return S_OK;
}

STDMETHODIMP CFSDrives::GetFolderProperty(PROPID propID, PROPVARIANT* value)
{
	try
	{
		CComPropVariant prop;
		switch(propID)
		{
			case kpidType:
				prop = CComBSTR(_T("FSDrives"));
				break;
			case kpidPath:
				if (_volumeMode)
				{
					prop = CComBSTR(kVolPrefix);
				}
				else
				{
					prop = CComBSTR(Lang::LangString(IDS_COMPUTER, 0x03020300) + CString(CHAR_PATH_SEPARATOR));
				}
				break;
		}
		prop.Detach(value);
		return S_OK;
	}
	catch(...)
	{
		return E_OUTOFMEMORY;
	}
}


STDMETHODIMP CFSDrives::GetSystemIconIndex(UINT32 index, INT32* iconIndex)
{
	*iconIndex = 0;
	const CDriveInfo &di = _drives[index];
	int iconIndexTemp;
	if (Common::FileSystem::GetRealIconIndex(di.FullSystemName, 0, iconIndexTemp) != 0)
	{
		*iconIndex = iconIndexTemp;
		return S_OK;
	}
	return GetLastError();
}

CString CFSDrives::GetExt(int index) const
{
	const CDriveInfo &di = _drives[index];
	const wchar_t* ext = NULL;
	if (di.DriveType == DRIVE_CDROM)
	{
		ext = _T("iso");
	}
	else if (di.FileSystemName.Find(_T("NTFS")) >= 0)
	{
		ext = _T("ntfs");
	}
	else if (di.FileSystemName.Find(_T("FAT")) >= 0)
	{
		ext = _T("fat");
	}
	else
	{
		ext = _T("img");
	}
	return (CString)_T('.') + ext;
}

HRESULT CFSDrives::GetLength(int index, UINT64 &length) const
{
	Common::IO::CInFile inFile;
	if (!inFile.Open(_drives[index].GetDeviceFileIoName()))
	{
		return GetLastError();
	}
	if (!inFile.LengthDefined)
	{
		return E_FAIL;
	}
	length = inFile.Length;
	return S_OK;
}

STDMETHODIMP CFSDrives::CopyTo(const UINT32* indices, UINT32 numItems,
							   const wchar_t* path, IFolderOperationsExtractCallback* callback)
{
	if (numItems == 0)
	{
		return S_OK;
	}

	if (!_volumeMode)
	{
		return E_NOTIMPL;
	}

	UINT64 totalSize = 0;
	UINT32 i;
	for (i = 0; i < numItems; i++)
	{
		const CDriveInfo &di = _drives[indices[i]];
		if (di.KnownSizes)
		{
			totalSize += di.DriveSize;
		}
	}
	HRESULT hRes(callback->SetTotal(totalSize));
	if(S_OK != hRes)
	{
		return hRes;
	}
	hRes = callback->SetNumFiles(numItems);
	if(S_OK != hRes)
	{
		return hRes;
	}

	CString destPath = path;
	if (destPath.IsEmpty())
	{
		return E_INVALIDARG;
	}
	bool directName = (destPath[destPath.GetLength() - 1] != CHAR_PATH_SEPARATOR);
	if (directName)
	{
		if (numItems > 1)
		{
			return E_INVALIDARG;
		}
	}

	UINT64 completedSize = 0;
	hRes = callback->SetCompleted(&completedSize);
	if(S_OK != hRes)
	{
		return hRes;
	}
	for (i = 0; i < numItems; i++)
	{
		int index = indices[i];
		const CDriveInfo &di = _drives[index];
		CString destPath2 = destPath;
		CString name = di.Name;
		if (!directName)
		{
			CString destName = name;
			if (!destName.IsEmpty() && destName[destName.GetLength() - 1] == _T(':'))
			{
				destName.Delete(destName.GetLength() - 1);
				destName += GetExt(index);
			}
			destPath2 += destName;
		}
		CString srcPath = di.GetDeviceFileIoName();

		UINT64 fileSize = 0;
		if (GetLength(index, fileSize) != S_OK)
		{
			return E_FAIL;
		}
		if (!di.KnownSizes)
		{
			totalSize += fileSize;
		}
		hRes = callback->SetTotal(totalSize);
		if(S_OK != hRes)
		{
			return hRes;
		}

		INT32 writeAskResult;
		CComBSTR destPathResult;
		hRes = callback->AskWrite(srcPath, BoolToInt(false), NULL, &fileSize,
								  destPath2, &destPathResult, &writeAskResult);
		if(S_OK != hRes)
		{
			return hRes;
		}
		if (!IntToBool(writeAskResult))
		{
			continue;
		}

		hRes = callback->SetCurrentFilePath(srcPath);
		if(S_OK != hRes)
		{
			return hRes;
		}

		static const UINT32 kBufferSize = (4 << 20);
		UINT32 bufferSize = (di.DriveType == DRIVE_REMOVABLE) ? (18 << 10) * 4 : kBufferSize;
		hRes = CopyFileSpec(srcPath, destPathResult, false, fileSize, bufferSize, completedSize, callback);
		if(S_OK != hRes)
		{
			return hRes;
		}
		completedSize += fileSize;
	}
	return S_OK;
}

STDMETHODIMP CFSDrives::MoveTo(
	const UINT32* /* indices */,
	UINT32 /* numItems */,
	const wchar_t* /* path */,
	IFolderOperationsExtractCallback* /* callback */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::CopyFrom(const wchar_t* /* fromFolderPath */,
								 const wchar_t** /* itemsPaths */, UINT32 /* numItems */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::CreateFolder(const wchar_t* /* name */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::CreateFile(const wchar_t* /* name */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::Rename(UINT32 /* index */, const wchar_t* /* newName */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::Delete(const UINT32* /* indices */, UINT32 /* numItems */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

STDMETHODIMP CFSDrives::SetProperty(UINT32 /* index */, PROPID /* propID */,
									const PROPVARIANT* /* value */, IProgress* /* progress */)
{
	return E_NOTIMPL;
}

}
}